﻿#include "rtmpose_utils.h"

RTMPose_Utils::RTMPose_Utils()
{
}

RTMPose_Utils::~RTMPose_Utils()
{
}

cv::Mat RTMPose_Utils::GetAffineTransform(float center_x, float center_y,
                                          float scale_width, float scale_height,
                                          int output_image_width, int output_image_height,
                                          bool inverse)

{
    // 获得原图与目标图的三个点
    cv::Point2f src_point_1;
    src_point_1.x = center_x;
    src_point_1.y = center_y;

    cv::Point2f src_point_2;
    src_point_2.x = center_x;
    src_point_2.y = center_y - scale_height * 0.5;

    cv::Point2f src_point_3;
    src_point_3.x = src_point_2.x - (src_point_1.y - src_point_2.y);
    src_point_3.y = src_point_2.y - (src_point_1.x - src_point_2.x);

    float alphapose_image_center_x = output_image_width / 2;
    float alphapose_image_center_y = output_image_height / 2;

    cv::Point2f dst_point_1;
    dst_point_1.x = alphapose_image_center_x;
    dst_point_1.y = alphapose_image_center_y;

    cv::Point2f dst_point_2;
    dst_point_2.x = alphapose_image_center_x;
    dst_point_2.y = alphapose_image_center_y - output_image_height * 0.5;

    cv::Point2f dst_point_3;
    dst_point_3.x = dst_point_2.x - (dst_point_1.y - dst_point_2.y);
    dst_point_3.y = dst_point_2.y - (dst_point_1.x - dst_point_2.x);

    cv::Point2f srcPoint[3] = {src_point_1, src_point_2, src_point_3};
    cv::Point2f dstPoint[3] = {dst_point_1, dst_point_2, dst_point_3};

    // 获取仿射变换矩阵 inverse=true为正向
    cv::Mat affinetransform;
    if (inverse)
    {
        affinetransform = cv::getAffineTransform(srcPoint, dstPoint);
    }
    else
    {
        affinetransform = cv::getAffineTransform(dstPoint, srcPoint);
    }

    return affinetransform;
}

float RTMPose_Utils::LetterBoxImage(const cv::Mat &image,
                                    cv::Mat &out_image,
                                    const cv::Size &new_shape,
                                    int stride,
                                    const cv::Scalar &color,
                                    bool fixed_shape,
                                    bool scale_up)
{
    cv::Size shape = image.size();
    float r = std::min((float)new_shape.height / (float)shape.height, (float)new_shape.width / (float)shape.width);

    if (!scale_up)
    {
        r = std::min(r, 1.0f);
    }

    int newUnpad[2]{
        (int)std::round((float)shape.width * r), (int)std::round((float)shape.height * r)};

    cv::Mat tmp;
    if (shape.width != newUnpad[0] || shape.height != newUnpad[1])
    {
        cv::resize(image, tmp, cv::Size(newUnpad[0], newUnpad[1]));
    }
    else
    {
        tmp = image.clone();
    }

    float dw = new_shape.width - newUnpad[0];
    float dh = new_shape.height - newUnpad[1];

    if (!fixed_shape)
    {
        dw = (float)((int)dw % stride);
        dh = (float)((int)dh % stride);
    }

    int top = int(0);
    int bottom = int(std::round(dh + 0.1f));
    int left = int(std::round(0));
    int right = int(std::round(dw + 0.1f));

    cv::copyMakeBorder(tmp, out_image, top, bottom, left, right, cv::BORDER_CONSTANT, color);

    return 1.0f / r;
}